# Bandits with dependent arms arms help to explore the space when
# we have to many arms to try all of them and can benefit because of 
# talking dependences into account

import numpy
import math
import random
import Distance
import os

class DS(object):
    # SIGMA for Gaussian kernel
    SIGMA = 0.3
    # NOISE for distance matrix
    NOISE = 0.05
    
    def __init__(self, file, images_number, setsize):
        # calculated mean
        self.evaluated_mean = dict()
        # images that were chosen by the user at all previouse iterations
        self.images = []
        self.setsize = setsize
        distance_type = 'Euclidean'
        self.images_number = images_number
        # file to get distance matrix
        self.path1 = os.path.dirname(__file__) + '/ProximityMatrix/'
        # from input file we will read matrix of distances
        # in the kernel version we had Gaussian kernel already computed
        # self.input = self.path1 + file + str(images_number) + '_kernel.dist'   
        self.input = self.path1 + file + str(images_number) + '_Euclidean.dist'   
        
        # previous_images contains images presented on the previuose iteration not to repeat the same image twice in a row
        self.previous_images = []
        self.LoadKernel()
        
        # sufficient statistics for decision making - mean and standard deviation
        self.setsize = 10
                        
        # variables for sampling
        self.training_points_kernel = numpy.array(())
        self.training_points_kernel_inversed = numpy.array(())
      
    # function for reading distances and transforming them to a kernel  
    def LoadKernel(self):
        # first load distances
        self.kernel = numpy.genfromtxt(self.input, unpack=True)
        # calculated Gaussian kernel
        self.kernel = numpy.exp(-(self.kernel)/(2*(DS.SIGMA**2)))
      
    # main function for predicting images  
    def Sampling(self, relevance):
        # add to self.images images that were chosen by the user
        for image in relevance:
            if relevance[image] == 1:
                self.images.append(image)
                
        # we will be keeping track of accepted images
        accepted_images = dict()
        # images to predict
        chosen_images = []
        # images that were not accepted but we may use them ig we run out of accepted ones
        refused_images = dict()
        
        # initialize accepted images with means predicted on previouse iterations or random numbers
        for image in self.images:
            accepted_images[image] = self.evaluated_mean.get(image, random.random()) 
        
        # for training_points_kernel choose lines corresponding to accepted images
        k = self.kernel[accepted_images.keys()]
        p = k.T
        self.training_points_kernel = p[accepted_images.keys()]
        # invert training_points
        self.training_points_kernel_inversed = numpy.linalg.inv(self.training_points_kernel)
        
        # for every image in the dataset
        for image in range(self.images_number):
            # if it was not presented on the previouse iteration
            if image not in relevance.keys():
                # HACK - and it was not chosen by the user before
                if image not in self.images:
                    # evaluate image's mean and variance
                    evaluated_variance, evaluated_mean = self.EvaluateMeanVariance(image, accepted_images)
                    # if evaluated variance bigger than the maximum that we have for accepted images add image to accepted set and recalculate differnt statistics
                    if evaluated_variance > max(accepted_images.values()):
                        accepted_images[image] = evaluated_variance
                        k = self.kernel[accepted_images.keys()]
                        p = k.T
                        self.training_points_kernel = (p[accepted_images.keys()])
                        self.training_points_kernel_inversed = numpy.linalg.inv(self.training_points_kernel)
                    else:
                        # if image was not selected add it to refused set
                        refused_images[image] = evaluated_variance
                 
        # take top accepted images   
        for key, value in sorted(accepted_images.iteritems(), key=lambda (k,v): (v,k), reverse=True):
            if len(chosen_images) < self.setsize:
                if key not in relevance.keys():
                    chosen_images.append(key)
            else:
                break
            
        # if there is not enought accepted images, take them from refused set
        for key, value in sorted(refused_images.iteritems(), key=lambda (k,v): (v,k), reverse=True):
            if len(chosen_images) < self.setsize:
                if key not in relevance.keys():
                    chosen_images.append(key)
            else:
                break
            
        return chosen_images
    
    # function for calculating mean and covariance of the datapoint      
    def EvaluateMeanVariance(self, image, accepted_images):
        k = self.kernel[accepted_images.keys()]
        image_repr = k.T[image]
        m = numpy.dot(numpy.dot(image_repr, self.training_points_kernel_inversed), numpy.array(accepted_images.values()).T)
        var = 1 - numpy.dot(numpy.dot(image_repr, self.training_points_kernel_inversed), image_repr.T)        
        evaluated_variance = m + math.sqrt(var)
        randn = random.normalvariate(0, 1)
        # sample from gaussian distribution
        self.evaluated_mean[image] = math.sqrt(var)*randn + m
        
        return evaluated_variance, self.evaluated_mean[image]
        
        
    # sample #setsize images randomly at the first iteration to present to the user and get first feedback
    def ChooseFirstImageSet(self):
        chosen_images = []
        while len(chosen_images) < self.setsize:
            image = random.randint(0, self.images_number-1)
            if image not in chosen_images:
                chosen_images.append(image)
        # save images as been presented
        self.previous_images = chosen_images
        # update timing
        self.T += self.setsize
        
        return chosen_images
